import torch

__all__ = ["EncoderLSTM"]


class EncoderLSTM(torch.nn.Module):
    def __init__(
        self, embedding: torch.nn.Embedding, hid_dim: int, n_layers: int, dropout: float
    ):
        super().__init__()
        self.hid_dim = hid_dim
        self.n_layers = n_layers
        self.embedding = embedding
        self.rnn = torch.nn.LSTM(
            embedding.embedding_dim,
            hid_dim,
            n_layers,
            dropout=dropout,
            batch_first=True,
        )
        self.dropout = torch.nn.Dropout(dropout)

    def forward(self, src):
        embedded = self.dropout(self.embedding(src))
        outputs, (hidden, cell) = self.rnn(embedded)
        return hidden, cell, outputs
